import easyocr
import numpy as np
from typing import List
import warnings
from .base import BaseOcrEngine
from ..function.recognize import OcrResult

# 忽略GPU相关警告
warnings.filterwarnings('ignore', category=UserWarning)

class EasyOcrEngine(BaseOcrEngine):
    def __init__(self, languages: List[str] = ['ch_sim', 'en']):
        self.reader = easyocr.Reader(languages, gpu=False)  # 强制使用CPU
        
    def recognize(self, image: np.ndarray, confidence_threshold: float = 0.5) -> List[OcrResult]:
        results = []
        ocr_results = self.reader.readtext(image)
        for box, text, confidence in ocr_results:
            if confidence >= confidence_threshold:
                results.append(OcrResult(text, confidence, box))
        return results 